#include "precompiled.h"

#include "CScanner.h"
#include "CUtils.h"

#include <windows.h>

std::map< std::string, std::string > CScanner::sign_base;

typedef std::map< std::string, std::string >::const_iterator sign_base_iter;

bool CScanner::readBase(const char* base)
{
	sign_base.clear();

	TiXmlDocument doc(base);
	if (!doc.LoadFile()) return false;

	TiXmlElement * first_sign = doc.RootElement()->FirstChildElement("signature");

	std::string sign_bin;
	do
	{
		CUtils::decodeBase64( first_sign->GetText(), sign_bin );
		sign_base[ first_sign->Attribute("name") ] = sign_bin;
	}
	while ( (first_sign = first_sign->NextSiblingElement("signature")) );

	return true;
}

bool CScanner::scan( const char* in, const char* base, const char* public_key )
{
	if ( !CUtils::validateSign(base, public_key) )
	{
		std::cout << "The signature mismatch." << std::endl;
		return false;
	}

	readBase(base);

	if ( (GetFileAttributes(in) & FILE_ATTRIBUTE_DIRECTORY) ) return scanDirectory(in);

	return scanFile(in);
}

bool CScanner::scanFile(const char* in)
{
	FILE *file = fopen(in, "rb");
	if (!file) return false;

	std::cout << "Scanning " << in << " .... " << std::endl;

	DWORD start = GetTickCount();
	bool res = scanFile(in, file);
	if (!res)
	{
		float time = ((float)GetTickCount() - (float)start)/1000.0f;
		std::cout << "Time:  " << time << " s" << std::endl;
		std::cout << "Speed: " << (float)CUtils::getFileSize(file)/time << " bytes/sec" << std::endl;
		std::cout << "Speed: " << (float)sign_base.size()/time << " sign/sec" << std::endl;
	}

	fclose(file);
	return res;
}

bool CScanner::scanFile(const char* in, FILE * file)
{
	sign_base_iter sign = sign_base.begin(), end = sign_base.end();
	for (; sign != end; ++sign)
	{
		if ( !has(file, sign->second) ) continue;

		std::cout << "The file " << in << " was infected by " << sign->first << std::endl;
		return true;
	}
	return false;
}

const unsigned char *CScanner::find(const unsigned char* haystack, size_t hlen, const unsigned char* needle, size_t nlen)
{
    size_t scan = 0;
    size_t bad_char_skip[UCHAR_MAX + 1];
 
    if (nlen <= 0 || !haystack || !needle) return NULL;
 
    for (scan = 0; scan <= UCHAR_MAX; scan = scan + 1)
	{
        bad_char_skip[scan] = nlen;
	}
 
    size_t last = nlen - 1;
    for (scan = 0; scan < last; scan = scan + 1)
	{
        bad_char_skip[needle[scan]] = last - scan;
	}

    while (hlen >= nlen)
    {
        for (scan = last-1; haystack[scan] == needle[scan]; scan = scan - 1)
		{
            if (scan == 0) return haystack;
		}

        hlen     -= bad_char_skip[haystack[last]];
        haystack += bad_char_skip[haystack[last]];
    }
 
    return NULL;
}

bool CScanner::scanDirectory(const char* in)
{
	WIN32_FIND_DATA FindFileData;

	std::string path = std::string(in) + "\\";
	std::string mask = path + "*";

	HANDLE hFile = FindFirstFile(mask.c_str(), &FindFileData);

	if (hFile == INVALID_HANDLE_VALUE) return false;

	bool res = false;
	BOOL done = FALSE;
	for (; !done; done = !FindNextFile(hFile,&FindFileData))
	{
		if ( !strcmp(FindFileData.cFileName, ".") ) continue;
		if ( !strcmp(FindFileData.cFileName, "..") ) continue;
		if ( !strcmp(FindFileData.cFileName, ".svn") ) continue;

		if ( (GetFileAttributes((path + FindFileData.cFileName).c_str()) & FILE_ATTRIBUTE_DIRECTORY) )
		{
			scanDirectory((path + FindFileData.cFileName).c_str());
		}
		else
		{
			scanFile( (path + FindFileData.cFileName).c_str() );
		}
	}

	FindClose(hFile);

	return res;
}

bool CScanner::has(FILE* file, const std::string& signature)
{
	fseek(file, 0, SEEK_SET);

	bool res = false;

	long wnd_size = 200 * signature.size();
	long step = wnd_size - signature.size() + 1;
	long pos = 0;

	char * ptr = NULL;
	char buffer[8192] = {0};
	if (wnd_size >= 8192 ) ptr = new char[wnd_size];
	else ptr = buffer;

	size_t bytes_read = 0;
	while( (bytes_read = fread(ptr, 1, wnd_size, file)) )
	{
		pos += step;
		fseek(file, pos, SEEK_SET);

		//std::string str(ptr, bytes_read);
		//if (str.find(signature) == std::string::npos) continue;
		if ( !find((unsigned char*)ptr, bytes_read, (unsigned char*)signature.c_str(), signature.size()) ) continue;

		res = true;
		break;
	}

	if (wnd_size >= 8192 ) delete [] ptr;
	return res;
}
